# coding=utf-8
import os
import sys

import fitz
import pandas
from rich.progress import track

# ORIGIN_DIR stores all the pdf files
ORIGIN_DIR = r'0428'
ANNO_DIR = ORIGIN_DIR + ' with annotations'


class PDFWordHighlighter:

    def __init__(self):
        self.wordcount = {}
        self.words = []
        self.pdf_list = []

    def get_and_make_all_dir(self, dir_name: str):
        self.wordcount[dir_name] = {}
        if not os.path.exists(ANNO_DIR):
            os.mkdir(ANNO_DIR)
        for root, ds, fs in os.walk(dir_name):
            for directory in ds:
                new_root = root.replace(ORIGIN_DIR, ANNO_DIR)
                if not os.path.exists(os.path.join(new_root, directory)):
                    os.mkdir(os.path.join(new_root, directory))
                self.wordcount[os.path.join(root, directory)] = {}
            for file in fs:
                if file.endswith('.pdf'):
                    self.pdf_list.append(os.path.join(root, file))
                    self.wordcount[os.path.join(root, file)] = {}


    def get_words(self, words_filename):
        with open(words_filename) as wordcount_txt:
            for one_word in wordcount_txt:
                self.words.append(one_word.strip('\n'))
        for key in self.wordcount.keys():
            for one_word in self.words:
                self.wordcount[key][one_word] = 0


    def highlight_pdf(self, clear_pdf=True, add_highlight=True):
        for pdf_file in track(self.pdf_list, description='Processing PDF...'):
            # open a pdf file in the original directory
            with fitz.Document(filename='\\\?\\' + sys.path[0] + '\\' + pdf_file,
                               filetype='pdf') as pdf:
                # check if it's a pdf or not, just to make sure
                if not pdf.is_pdf:
                    continue

                for page in pdf:
                    # if clear_pdf is True, clear all the pdf annotations first
                    if clear_pdf:
                        for annot in page.annots():
                            page.delete_annot(annot)

                    for one_word in self.words:
                        # find the keywords in one page
                        word_num = page.get_textpage().search(one_word)
                        self.wordcount[pdf_file][one_word] += word_num.__len__()

                        # calculate the words for one directory
                        dirname = os.path.dirname(pdf_file)
                        while dirname != os.path.dirname(dirname):
                            self.wordcount[dirname][one_word] += word_num.__len__()
                            dirname = os.path.dirname(dirname)

                        # highlight all the keywords mentioned in the wordcount
                        for word_pos in word_num:
                            if add_highlight:
                                page.add_highlight_annot(word_pos)
                            pass

                # Save pdf as last
                pdf.save('\\\?\\' + sys.path[0] + '\\' + pdf_file.replace(ORIGIN_DIR, ANNO_DIR))


    def write_csv(self, filename=f'{ANNO_DIR}.csv'):
        df = pandas.DataFrame(self.wordcount)
        df.to_csv(filename, index=True, sep=',')


if __name__ == "__main__":
    highlighter = PDFWordHighlighter()
    highlighter.get_and_make_all_dir(ORIGIN_DIR)
    highlighter.get_words(words_filename='words.txt')
    highlighter.highlight_pdf(add_highlight=True, clear_pdf=True)
    highlighter.write_csv()
